from sklearn.cluster import KMeans
import csv
import math
import numpy as np
import matplotlib.pyplot as plt

#读取csv
filename = 'homework_03_kmeans\dataset_circles.csv'
X_point = []
Y_point = []
types = []
with open(filename) as f:
    reader = csv.reader(f)
    # 读取一行，下面的reader中已经没有该行了
    head_row = next(reader)
    X_point.append(head_row[0])
    Y_point.append(head_row[1])
    types.append(head_row[2])
    for row in reader:
        # 行号从2开始
        X_point.append(row[0])
        Y_point.append(row[1])
        types.append(row[2])
x = np.array(X_point, dtype=float)
y = np.array(Y_point, dtype=float)
point_type = np.array(types, dtype=float)

# #test print
# plt.scatter(x,y)
# plt.show()

# Feature transformation

def distanceArray(array_x, array_y):
    distance = []
    for i in range(len(array_x)):
        num = array_x[i]*array_x[i]+array_y[i]*array_y[i]
        dis = math.sqrt(num)
        distance.append(dis)
        distanceList = np.array(distance)
    return (distanceList)

distance_Array = distanceArray(x,y)
zeros = np.zeros(len(distance_Array))
newData = distance_Array[:, np.newaxis]

# K means
cluster_pred = KMeans(n_clusters=2, random_state=9).fit_predict(newData)

# Print result
for i in range(len(x)):
    if cluster_pred[i] == 1:
        plt.scatter(x[i], y[i], c='b')
    else:
        plt.scatter(x[i], y[i], c='r')

plt.show()
